import numpy as np
from torchvision.models import resnet50
from thop import profile
import torch
from model.edge_model import SegUNet
from thop import clever_format
from torch.nn import functional as F
model = SegUNet(9)
input = torch.randn(1, 3, 960, 544)
macs, params = profile(model, inputs=(input, ))
print(macs, params)
macs, params = clever_format([macs, params], "%.3f")
print(macs, params)

a = np.asarray(
    [
        [-1000, 1000],
        [-10.0, 10]
    ]
)
print(F.softmax(torch.from_numpy(a), dim=1))